import os
import json
import time
import queue
import shutil
import threading
from pathlib import Path
from tkinter import Tk, ttk, StringVar, filedialog, messagebox

from decrypt import decrypt_database_file
from dbutils import get_sessions, get_messages_for_session


APP_DIR = Path.home() / "Documents" / "EchoTrace"
CONFIG_PATH = APP_DIR / "config.json"


def load_config() -> dict:
    if CONFIG_PATH.exists():
        try:
            return json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
        except Exception:
            return {}
    return {}


def save_config(cfg: dict) -> None:
    APP_DIR.mkdir(parents=True, exist_ok=True)
    CONFIG_PATH.write_text(json.dumps(cfg, ensure_ascii=False, indent=2), encoding="utf-8")


def auto_detect_root() -> str | None:
    home = os.environ.get("USERPROFILE") or os.environ.get("HOME")
    if not home:
        return None
    p = Path(home) / "Documents" / "xwechat_files"
    return str(p) if p.exists() else None


class EchotraceApp:
    def __init__(self, root: Tk) -> None:
        self.root = root
        self.root.title("EchoTrace 参考实现 (Python)")
        self.root.geometry("1024x700")

        self.cfg = load_config()
        self.root_path = StringVar(value=self.cfg.get("root_path", ""))
        self.decrypt_key = StringVar(value=self.cfg.get("decrypt_key", ""))
        self.selected_account = StringVar()

        self.db_files = []  # list of dict
        self.sessions = []

        nb = ttk.Notebook(root)
        nb.pack(fill="both", expand=True)

        self.page_settings = ttk.Frame(nb)
        self.page_data = ttk.Frame(nb)
        self.page_chat = ttk.Frame(nb)
        nb.add(self.page_settings, text="设置")
        nb.add(self.page_data, text="数据管理")
        nb.add(self.page_chat, text="聊天记录")

        self._build_settings()
        self._build_data()
        self._build_chat()

        # 初始扫描
        self._scan_databases()

    # ---------- 设置 ----------
    def _build_settings(self) -> None:
        frame = self.page_settings
        pad = {"padx": 8, "pady": 6}

        row = 0
        ttk.Label(frame, text="数据库根目录 (xwechat_files)").grid(row=row, column=0, sticky="w", **pad)
        ttk.Entry(frame, textvariable=self.root_path, width=70).grid(row=row, column=1, sticky="we", **pad)
        ttk.Button(frame, text="浏览...", command=self._choose_root).grid(row=row, column=2, **pad)
        ttk.Button(frame, text="自动检测", command=self._auto_detect).grid(row=row, column=3, **pad)
        row += 1

        ttk.Label(frame, text="解密密钥(64位hex)").grid(row=row, column=0, sticky="w", **pad)
        ttk.Entry(frame, textvariable=self.decrypt_key, width=70, show="*").grid(row=row, column=1, sticky="we", **pad)
        ttk.Button(frame, text="保存", command=self._save_settings).grid(row=row, column=2, **pad)
        frame.grid_columnconfigure(1, weight=1)

    def _choose_root(self) -> None:
        d = filedialog.askdirectory(title="选择 xwechat_files 根目录")
        if d:
            self.root_path.set(d)

    def _auto_detect(self) -> None:
        d = auto_detect_root()
        if d:
            self.root_path.set(d)
            messagebox.showinfo("提示", f"自动检测成功: {d}")
        else:
            messagebox.showwarning("提示", "未检测到 xwechat_files 目录")

    def _save_settings(self) -> None:
        cfg = {
            "root_path": self.root_path.get().strip(),
            "decrypt_key": self.decrypt_key.get().strip(),
        }
        save_config(cfg)
        messagebox.showinfo("已保存", "设置已保存")

    # ---------- 数据管理 ----------
    def _build_data(self) -> None:
        frame = self.page_data
        toolbar = ttk.Frame(frame)
        toolbar.pack(fill="x")
        ttk.Button(toolbar, text="刷新", command=self._scan_databases).pack(side="left", padx=4, pady=4)
        ttk.Button(toolbar, text="批量解密未解密", command=self._batch_decrypt).pack(side="left", padx=4, pady=4)
        ttk.Button(toolbar, text="增量更新有变更", command=self._incremental_update).pack(side="left", padx=4, pady=4)

        columns = ("account", "file", "size", "src_mtime", "dec_exists", "dec_mtime", "needs_update")
        self.tree = ttk.Treeview(frame, columns=columns, show="headings")
        for c, w in zip(columns, (120, 200, 80, 160, 90, 160, 100)):
            self.tree.heading(c, text=c)
            self.tree.column(c, width=w, anchor="w")
        self.tree.pack(fill="both", expand=True)

        self.status_var = StringVar()
        ttk.Label(frame, textvariable=self.status_var).pack(anchor="w", padx=6, pady=4)

    def _scan_databases(self) -> None:
        root = Path(self.root_path.get().strip())
        self.db_files = []
        self.tree.delete(*self.tree.get_children())
        if not root.exists():
            self.status_var.set("未找到根目录")
            return
        total = 0
        for account_dir in root.iterdir():
            if not account_dir.is_dir():
                continue
            db_storage = account_dir / "db_storage"
            if not db_storage.exists():
                continue
            account_name = account_dir.name
            for p in db_storage.rglob("*.db"):
                total += 1
                size = p.stat().st_size
                src_mtime = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(p.stat().st_mtime))
                # 解密输出路径
                out_dir = APP_DIR / account_name
                out_path = out_dir / (p.name.split(".")[0] + ".db")
                dec_exists = out_path.exists()
                dec_mtime = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(out_path.stat().st_mtime)) if dec_exists else ""
                needs_update = (not dec_exists) or (out_path.stat().st_mtime < p.stat().st_mtime if dec_exists else True)
                info = {
                    "account": account_name,
                    "src": str(p),
                    "dst": str(out_path),
                    "needs_update": needs_update,
                }
                self.db_files.append(info)
                self.tree.insert("", "end", values=(account_name, p.name, f"{size//1024} KB", src_mtime, str(dec_exists), dec_mtime, "是" if needs_update else "否"))
        self.status_var.set(f"共发现 {total} 个数据库文件")

    def _batch_decrypt(self) -> None:
        key = self.decrypt_key.get().strip()
        if len(key) != 64:
            messagebox.showerror("错误", "请先在设置中填写 64 位十六进制密钥")
            return
        tasks = [f for f in self.db_files if not Path(f["dst"]).exists()]
        if not tasks:
            messagebox.showinfo("提示", "没有需要解密的文件")
            return
        threading.Thread(target=self._run_decrypt_tasks, args=(tasks, key), daemon=True).start()

    def _incremental_update(self) -> None:
        key = self.decrypt_key.get().strip()
        if len(key) != 64:
            messagebox.showerror("错误", "请先在设置中填写 64 位十六进制密钥")
            return
        tasks = [f for f in self.db_files if f["needs_update"]]
        if not tasks:
            messagebox.showinfo("提示", "没有需要更新的文件")
            return
        threading.Thread(target=self._run_decrypt_tasks, args=(tasks, key), daemon=True).start()

    def _run_decrypt_tasks(self, tasks: list[dict], key: str) -> None:
        ok = 0
        fail = 0
        for item in tasks:
            src = Path(item["src"]) ; dst = Path(item["dst"]) ; dst.parent.mkdir(parents=True, exist_ok=True)
            tmp = dst.with_suffix(".tmp.db")
            try:
                decrypt_database_file(str(src), key, str(tmp))
                # 替换
                if dst.exists():
                    try:
                        dst.unlink()
                    except Exception:
                        pass
                shutil.move(str(tmp), str(dst))
                ok += 1
            except Exception as e:
                fail += 1
            self._scan_databases()
        # 解密完成后刷新账号列表（在主线程刷新UI）
        self.root.after(0, self._refresh_accounts)
        messagebox.showinfo("完成", f"解密完成 成功:{ok} 失败:{fail}")

    # ---------- 聊天记录 ----------
    def _build_chat(self) -> None:
        frame = self.page_chat
        top = ttk.Frame(frame)
        top.pack(fill="x")
        ttk.Label(top, text="选择账号:").pack(side="left", padx=4, pady=4)
        self.cmb_account = ttk.Combobox(top, textvariable=self.selected_account, state="readonly")
        self.cmb_account.pack(side="left", padx=4)
        ttk.Button(top, text="刷新账号", command=self._refresh_accounts).pack(side="left", padx=6)
        ttk.Button(top, text="加载会话", command=self._load_sessions).pack(side="left", padx=6)

        mid = ttk.Frame(frame)
        mid.pack(fill="both", expand=True)
        self.sessions_tree = ttk.Treeview(mid, columns=("username", "display"), show="headings")
        self.sessions_tree.heading("username", text="username")
        self.sessions_tree.heading("display", text="display")
        self.sessions_tree.column("username", width=240)
        self.sessions_tree.column("display", width=240)
        self.sessions_tree.pack(side="left", fill="y")
        self.sessions_tree.bind("<<TreeviewSelect>>", self._on_select_session)

        self.msg_tree = ttk.Treeview(mid, columns=("time", "content"), show="headings")
        self.msg_tree.heading("time", text="time")
        self.msg_tree.heading("content", text="content")
        self.msg_tree.column("time", width=160)
        self.msg_tree.column("content", width=600)
        self.msg_tree.pack(side="left", fill="both", expand=True)

        # 初始化账号下拉
        self._refresh_accounts()

    def _accounts_dir(self) -> list[str]:
        """
        返回已解密可用的账号列表：Documents/EchoTrace/<wxid> 下存在 session.db 的目录。
        这样聊天页只显示“可读取聊天记录”的账号，避免未解密时列表为空的困惑。
        """
        if not APP_DIR.exists():
            return []
        return [d.name for d in APP_DIR.iterdir() if (d / "session.db").exists()]

    def _refresh_accounts(self) -> None:
        accounts = self._accounts_dir()
        self.cmb_account["values"] = accounts
        if accounts:
            self.selected_account.set(accounts[0])

    def _load_sessions(self) -> None:
        self.sessions_tree.delete(*self.sessions_tree.get_children())
        account = self.selected_account.get()
        if not account:
            return
        session_db = APP_DIR / account / "session.db"
        if not session_db.exists():
            messagebox.showwarning("提示", f"未找到解密的会话数据库: {session_db}")
            return
        self.sessions = get_sessions(str(session_db))
        for u, d, s in self.sessions:
            self.sessions_tree.insert("", "end", values=(u, d or ""))

    def _on_select_session(self, event=None) -> None:
        sel = self.sessions_tree.selection()
        if not sel:
            return
        username = self.sessions_tree.item(sel[0], "values")[0]
        account = self.selected_account.get()
        # 在 EchoTrace 下找到 message_*.db，逐个尝试
        account_dir = APP_DIR / account
        msg_dbs = sorted(account_dir.glob("message_*.db"))
        self.msg_tree.delete(*self.msg_tree.get_children())
        for dbp in msg_dbs:
            rows = get_messages_for_session(str(dbp), username, limit=100, offset=0)
            for r in rows:
                ts = r.get("create_time") or r.get("createTime") or 0
                if isinstance(ts, int):
                    from datetime import datetime
                    tstr = datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M:%S")
                else:
                    tstr = str(ts)
                content = r.get("message_content") or r.get("content") or ""
                self.msg_tree.insert("", "end", values=(tstr, content))


def main() -> None:
    root = Tk()
    app = EchotraceApp(root)
    root.mainloop()


if __name__ == "__main__":
    main()


